import pprint, copy, cPickle, gzip, tarfile, os, sys, pdb
import numpy as np
from pywolb.utilities import numpytools, stringtools
#from pywolb.utilities import animationtools
#import matplotlib.pyplot as plt
#from pywolb.deterministic import weights

for module in [numpytools,stringtools]:
    reload(module)


# optimizing for speed
expand = numpytools.expand
condense = numpytools.condense
normalize = numpytools.normalize

fa = numpytools.random_frequencies((2,3,2,2,2))
fl = dict( axes=['sex','background','trait','preference','cytotype'],\
           elements=[['female','male'],['S1','S2','SH'],['T1','T2'],['P0','P1'],['uninfected','Wolbachia']] )
cia = [[1.,1.],[0.1,1.]]
cil = dict( axes=['male cytotype', 'offspring cytotype'],\
            elements=[['mU','mW'],['oU','oW']] )
                
def reproduction_axes(*locus_names):
    axes = ["%s %s" % (i, name) for i in ['female','male','offspring'] for name in locus_names]
    return axes
    

def list2commastring(l):
    return '\'' + '\', \''.join(l) + '\''

class ParameterError(Exception):
    def __init__(self, func, expected, given):
        self.func = func
        self.expected = expected
        self.given = given
    
    def __str__(self):
        return repr(self)
    
    def __repr__(self):
        return """wrong argument(s) given to %s();
    expected: %s
    given: %s""" % (self.func, list2commastring(self.expected), list2commastring(self.given))        

class LabeledArray: #(object):
    # in childclasses, fill in modules and classes that the child class objects need to be instantiated:
    # e.g.: _prereqs= {'numpy': ['array'], 'pprint': ['pprint','pformat']}
    _return_class = 'LabeledArray'
    _sums = False        # set this to True if it makes sense to calculate sums along axes
    _string_threshold = 1000
    
    def __init__(self, data, labels=None):
        data = np.array(data)
        self.data = np.array(data)
        ds = self.data.shape
        self._data_shape = ds
        if labels == None:
            labels = self._default_labels(ds)  
        self._init_labels(labels)
        self._ndim = self.data.ndim
    
    def _default_labels(self, _shape):
        """If no labels are provided then create default labels.
        
        E.g.: If _data = np.array([[1.,2.],[3.,4.],[5.,6.]] then
              labels = {'axes': ['a0', 'a1'],
                        'elements': [['e00','e01','e02'],['e10','e11']]}.
              print a would then yield:
              [[1., 2.],   |   [['e00-e10', 'e00...
               [3., 4.],   |
               [5., 6.]]   |
        """
        n = len(_shape)
        axes = []
        elements = []
        for i in range(n):
            axes.append("ax_%d" % i)
            temp_list = []
            for j in range(_shape[i]):
                temp_list.append("el_%d%d" % (i,j))
            elements.append(temp_list)
        return dict(axes=axes,elements=elements)
    
    def _init_labels(self, labels):
        """Initialize label members.
        
        Label members are:
            labels: a copy of the labels dictionary,
            _get_axis_position: a dictionary with axis labels as keys and axis position as values,
            _get_element_postion: a dictionary with element labels as keys and element position as values,
            _types: a numpy array containing type labels.
        """
        self.labels = labels
        self._label_shape = numpytools.list_shape( labels['elements'] )
        self._axes = labels['axes']
        self._elements = numpytools.flatten( labels['elements'] )
        self._get_axis_position = self._create_axis_map( labels['axes'] )
        self._get_element_position = self._create_element_map( labels['elements'] )
        self._types = self._create_type_labels( labels['elements'] )
        self._get_type_position = self._create_type_map( self._types )
    
    def get_shape(self):
        return self._data_shape
    
    def _reinit(self):
        """LabeledArray must be reinitialized after loading from file."""
        self.__init__(data=self.data, labels=self.labels)
    
    def __array__(self):
        return self.data
    
    def __str__(self):
        return self._data_labels2string()
    
    #def __repr__(self):
        ##return "LabeledArray(%s, %s)" % (self.data, self.labels)
        #return self._data_labels2string()
    
    def _data_labels2string(self):
        th = self._string_threshold
        el = 5
        symb = (self._ndim-1)*' ' + '...'
        s1 = stringtools.string_summary( stringtools.strip_blank_lines( np.array2string(self.data) ), \
                                threshold=th, \
                                edgelines=el, \
                                sum_symbol=symb)
        s2 = stringtools.string_summary( stringtools.strip_blank_lines( np.array2string(self._types, max_line_width=1000) ), \
                                threshold=th, \
                                edgelines=el, \
                                sum_symbol=symb)
        # replace single quotes by spaces:
        s2 = s2.replace("'", " ")
        return stringtools.two_columns( s1, s2, num_spaces=6, sep_symbol= '')
    
    def _show(self, show_sums=False, show_entire=False):
        """Prints a pretty formatted presentation of the instance."""
        s = 'array:\n'
        if show_entire:
            threshold = self._string_threshold
            self._string_threshold = np.nan
            s += stringtools.indent( self._data_labels2string(), 4 ) + '\n'
            self._string_threshold = threshold
        else:
            s += stringtools.indent( self._data_labels2string(), 4 ) + '\n'
        s += 'axes:\n'
        for i in range(self._ndim):
            ax = self.labels['axes'][i]
            if self._sums and show_sums:
                s += stringtools.indent( "%s  %s:" % (i, ax), 4 ) + '\n'
                s += stringtools.indent( self._sums2string(ax), 12 ) + '\n'
            else:
                s += stringtools.indent( "%s  %s" % (i, ax), 4 ) + '\n'
        return s
    
    def show(self, show_sums=False, show_entire=False):
        print self._show(show_sums, show_entire)
    
    def _overview2string(self):
        s = ''
        for ax in self._axes:
            s += self._sums2string(ax) + '\n\n'
        return s
    
    def show_overview(self):
        print self._overview2string()
    
    def _sums2string(self, axis):
        out = []
        elements = []
        sums = []
        for el in self._get_axis_elements(axis):
            elements.append(el)
            sums.append(self.get_sum(el))
        longest = stringtools.max_length(elements)
        for i in range(len(elements)):
            code = "%%-%ds  %%s" % longest
            out.append( code % (elements[i], sums[i]) )
        return '\n'.join(out)
    
    
    def __getitem__(self, y):
        return self.get_slice(y)
    
    def _create_axis_map(self, axis_labels):
        _map = {}
        for axis in axis_labels:
            _map[axis] = axis_labels.index(axis)
        return _map
        
    def _create_element_map(self, element_labels):
        _map = {}
        for axis in element_labels:
            for element in axis:
                _map[element] = [ element_labels.index(axis), axis.index(element) ]
        return _map
    
    def _create_type_map(self, type_array):
        _map = {}
        if type_array.ndim != 0:
            for _type in numpytools.flatten(type_array):
                _map[_type] = self._get_index(_type)
        return _map
    
    def _get_index(self, _type):
        return numpytools.array_index(self._types, _type)
        
    def _create_type_labels(self, element_labels):
        plabels = numpytools.permute(element_labels)
        a = np.array( ['-'.join(row) for row in plabels] )
        if a.ndim == 0:
            return np.array(None)
        a.shape = self._label_shape
        return a
    
    def get_position(self, x):
        if x in self._axes:
            return self._get_axis_position[x]
        elif x in self._elements:
            return self._get_element_position[x]
        else:
            raise ValueError, "key '%s' not found" % x
    
    def get_label(self, index):
        return self._types[index]
    
    def _get_axis_elements(self, _axis):
        apos = self._get_axis_position[_axis]
        return self.labels['elements'][apos]
    
    def _get_element_axis(self, element):
        return self._axes[self.get_position(element)[0]]
    
    def get_slice(self, y):
        # allow normal slicing notation (but this will return a normal numpy array):
        try:
            if isinstance(y, (int,tuple,slice)):
                return self.data[y]
            elif isinstance(y, str):
                if y in self._axes:
                    return self._get_axis_slice(y)
                return self._get_type_slice(y)
        except ValueError:
            raise ValueError, "'%s' is not a valid input. Input must either be an integer, a slice object, a tuple of integers and/or slices, or a type string." % str(y)
    
    def _get_axis_slice(self, _axis):
        pos = self.get_position(_axis)
        result = np.sum(self.data, pos)
        newlabels = copy.deepcopy(self.labels)
        for x in ['axes', 'elements']:
            del newlabels[x][pos]
        code = "%s(data=result, labels=newlabels)" % self._return_class
        return eval(code)
        
    def _get_type_slice(self, _type):
        axis_labels = copy.copy( self.labels['axes'] )
        element_labels = copy.copy( self.labels['elements'] )
        positions = []
        for element in _type.split('-'):
            positions.append(self.get_position(element))
        slice_list = [slice(None,None,None) for i in range(self._ndim)]
        for ax,el in sorted(positions):
            slice_list[ax] = el
            ax_name = self.labels['axes'][ax]
            pos = axis_labels.index(ax_name)
            del axis_labels[pos]
            del element_labels[pos]
        result = self.data[ tuple(slice_list) ]
        if isinstance(result, float):
            return result
        labels = {'axes': axis_labels, 'elements': element_labels}
        code = "%s(data=result, labels=labels)" % self._return_class
        return eval(code)
    
    def get_sum(self, *y):
        if y == ():
            return np.sum(self)
        y = y[0]
        #arg = y.split('-')[0]        # check whether input is axis or type
        if y in self._axes:
            return self._get_axis_sum(y)
        try:
            return self._get_type_sum(y)
        except ValueError:
            raise ValueError, "'%s' is not a valid input. Input must either be an axis or type label." % str(y)
    
    def _get_axis_sum(self, _axis):
        #result = condense(self.data, self.get_position(_axis))
        #newlabels = dict(axes=[_axis], elements=[self._get_axis_elements(_axis)])
        #code = "%s(data=result, labels=newlabels)" % self._return_class
        #return eval(code)
        axes = {}        # { position: (axis_name, elements) }
        for axis in _axis.split('--'):
            axes[self.get_position(axis)] = ( axis, self._get_axis_elements(axis) )
        positions = sorted(axes)
        axis_names = []
        elements = []
        for p in positions:
            an, el = axes[p]
            axis_names.append(an)
            elements.append(el)
        result = condense(self.data, positions)
        newlabels = dict(axes=axis_names, elements=elements)
        code = "%s(data=result, labels=newlabels)" % self._return_class
        return eval(code)
    
    def _get_type_sum(self, _type):
        """Example _type: 'Wolbachia-T1-P0'"""
        positions = []
        for element in _type.split('-'):
            positions.append(self.get_position(element))
        axis_positions = []
        element_positions = []
        for ax,el in sorted(positions):
            axis_positions.append(ax)
            element_positions.append(el)
        return condense(self.data, axis_positions)[tuple(element_positions)]
    
    def normalize(self, norm=1.):
        self.data *= norm/np.sum(self.data)
    
    def set_array(self, data):
#        data = copy.deepcopy( np.array(data) )
        data = np.array(data)
        if data.shape == self._data_shape:
            self.data = data[:]
        else:
            raise ValueError, "shape must be %s, instead %s given" % (str(self._data_shape), str(data.shape))
    
    # some method shortcuts / abbreviations:
    index = _get_index
    getl = get_label
    gets = getsum = get_sum

class Weight(LabeledArray):
    _return_class = 'LabeledArray'
    _type = 'base weight'
    _parameter_names = []
    _sums = False
    labels = None
    
    def __init__(self, **kwargs):
        """Possible kwargs:
_scope = 'global' / 'population' --- 'global' weights are the same in all populations, 'population' weights might differ between populations
parameters should also be provided by kwargs"""
        try:
            self.info = self.__class__._info
        except:
            pass
        self._scope = 'global'         # this gets overridden by child classes
        self.__dict__.update(kwargs)
        #for pname in parameter_names:
            #self.__dict__[pname] = parargs[pname]
        data = self._build()
        LabeledArray.__init__(self, data, self.labels)
            
    def _init_after_assignment(self):
        pass
    
    def _refer_to_pop(self, pop):
        """Store a reference to the population the weight belongs to.
        
        A population should call this method when a weight with scope 'population' is assigned to it."""
        self.pop = pop
    
    def _build(self):
        """This must return an array.
        
        E.g.:
        If parameters={'ci_level': 0.9} then
        _init_parameters creates a member variable self.ci_level = 0.9
        and a list self._properties['parargs'] = ['ci_level'].
        Thus, _build() could return np.array([[1., 1.], [1., 1-ci]]) 
        where ci = self.ci_level.
        """
        return np.ones((2,2))
    
    def _rebuild(self):
        self.set_array( self._build() )
    
    def _calc(self):
        """If the weight is a dynamic one then its data has to be calculated anew
        each generation. Use this method."""
        return np.ones((2,2))
    
    def set(self, **kwargs):
        parameter_names = kwargs.keys()
        for pname in parameter_names:
            if pname not in self._parameter_names:
                raise NameError, "parameter '%s' is not valid" % pname
        self.__dict__.update(kwargs)
        data = self._rebuild()
    
    def get_parnames(self):
        return self._parameter_names
    
    def update(self):
        self.data = self._calc()
        
    def _show(self, show_sums=False, show_entire=False):
        s = "weight for %s:\n" % self.info
        scope = self._scope
        if scope == 'global':
            s += stringtools.indent('scope: %s' % scope, 4) +'\n'
        elif scope == 'population':
            if self.pop:
                s += stringtools.indent("scope: population %s" % repr(self.pop), 4) +'\n'
            else:
                s += stringtools.indent('scope: %s (not bound yet)' % scope, 4) +'\n'
        npars = len(self._parameter_names)
        if npars >= 1:
            if npars == 1:
                s += stringtools.indent( 'parameter:', 4 ) + '\n'
            else:
                s += stringtools.indent( 'parameters:', 4 ) + '\n'
            for name in self._parameter_names:
#                s += stringtools.indent( "%s  %s" % (name, pprint.pformat(self.__dict__[name])), 8 ) + '\n'
                prettyparstring = pprint.pformat(self.__dict__[name]).replace('{',' ').replace('}','').replace(',','')
                s += stringtools.indent( stringtools.two_columns(name, prettyparstring), 8) + '\n'
        s += stringtools.indent( 'array:', 4 ) + '\n'
        if show_entire:
            threshold = self._string_threshold
            self._string_threshold = np.nan
            s += stringtools.indent( self._data_labels2string(), 8 ) + '\n'
            self._string_threshold = threshold
        else:
            s += stringtools.indent( self._data_labels2string(), 8 ) + '\n'
        s += stringtools.indent( 'axes:', 4) + '\n'
        for i in range(self._ndim):
            s += stringtools.indent( "%s  %s" % (i, self.labels['axes'][i]), 8 ) + '\n'
        return s
    
    # summing does not make sense in weighting factors:
    def get_sum(self, y):
        raise AttributeError, "Weight instances do not have sum methods."
    def _get_axis_sum(self, y):
        raise AttributeError, "Weight instances do not have sum methods."
    def _get_type_sum(self, y):
        raise AttributeError, "Weight instances do not have sum methods."
        
    def expanded(self, expanded_axes):
        new_ndim = len( expanded_axes )
        own_positions = [expanded_axes.index(_axis) for _axis in self._axes]
        return numpytools.expand( self.data, new_ndim, own_positions )
    
class Frequencies(LabeledArray):
    _return_class = 'Frequencies'
    _sums = True
    
    def _show(self, show_sums=True, show_entire=False, show_most_frequent=True):
        """Prints a pretty formatted presentation of the instance."""
        s = LabeledArray._show(self, show_sums, show_entire)
        if show_most_frequent:
            try:
                s += "most frequent types:\n"
                s += stringtools.indent( self._mostfrequent2string(8), 4 )
            except:
                pass
        return s
    
    def show(self, show_sums=True, show_entire=False, show_most_frequent=True):
        print self._show(show_sums, show_entire, show_most_frequent)
    
    def get_most_frequent(self, sum_axes=None, cutoff=1e-10):
        """Returns a list of tuples containing the type label and the according frequency,
sorted from most frequent to least frequent.
sum_axes might be a list of indexes of axes that should be summed along before the sorting procedure."""
        result = []
        code = "freqs = %s(data=self.data,labels=self.labels)" % self._return_class
        exec code
        if sum_axes is not None:
            if isinstance(sum_axes, str):
                sum_axes = [sum_axes]
            for axis in sum_axes:
                freqs = freqs[axis]
        sorted_indexes = numpytools.argsort(freqs.data)
        for i in sorted_indexes:
            if freqs[i] > cutoff:
                result.append( (freqs._types[i], freqs[i]) )
        return result[::-1]
    
    def _mostfrequent2string(self, n=5, sum_axes=None):
        out = []
        labels = []
        sums = []
        for l,s in self.get_most_frequent(sum_axes)[:n]:
            labels.append(l)
            sums.append(s)
        longest = stringtools.max_length(labels)
        for i in range(len(labels)):
            code = "%%-%ds  %%s" % longest
            out.append( code % (labels[i], sums[i]) )
        return '\n'.join(out)
    
    def _full_overview2string(self, n_most_frequent=5):
        s = ''
        for ax in self._axes:
            s += self._sums2string(ax) + '\n\n'
        s += 'most frequent:\n'
        s += self._mostfrequent2string(n_most_frequent)
        return s
    
    def show_full_overview(self, n_most_frequent=5):
        print self._full_overview2string(n_most_frequent)
    
class FrequenciesHistory(LabeledArray):
    _return_class = ' FrequenciesHistory'
    
    def __init__(self, generations=0, data=0., labels=None):
        # one or more generations provided?
        if isinstance(generations, int):
            case = 'single'
        if isinstance(generations, list):
            case = 'multiple'
            
        if case == 'single':
            generations = [generations]
            # make sure that data_list is a list of arrays:
#            data_list = [ copy.deepcopy( np.array(data)) ]
            data_list = [ np.array(data) ]
            # get labels from data if none are provided:
            if not labels:
                try:
                    labels = data.labels
                except:
                    labels = Frequencies(data).labels
        if case == 'multiple':
            generations = list(generations)
            # make sure that data_list is a list of arrays:
            data_list = [ np.array(d) for d in data ]
            if len(generations) != len(data_list):
                raise ValueError, "'generations' and 'data' do not have the same length"
            # get labels from data if none are provided:
            if not labels:
                try:
                    labels = data[0].labels
                except:
                    labels = Frequencies(data[0]).labels
        self._data_list = data_list
        self._freq_shape = np.shape(data_list[0])
        LabeledArray.__init__(self, data=self._data_list, labels=labels)
        self.generations = generations
        self._array_uptodate = True
    
    def get_slice(self, y):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        return LabeledArray.get_slice(self, y)
    
    def __array__(self):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        return self.data
    
    #def __repr__(self):
        #th = self._string_threshold
        #el = 5
        #symb = (self._ndim-1)*' ' + '...'
        #return stringtools.string_summary( stringtools.strip_blank_lines( np.array2string(self.data) ), \
                                  #threshold=th, \
                                  #edgelines=el, \
                                  #sum_symbol=symb)
    
    def __len__(self):
        return len(self.generations)
    
    def show(self, show_entire=False):
        print self._show(show_entire=show_entire)
    
    def _data_labels2string(self):
        th = self._string_threshold
        el = 5
        symb = (self._ndim-1)*' ' + '...'
        s1 = stringtools.string_summary( stringtools.strip_blank_lines( np.array2string(self.data) ), \
                                threshold=th, \
                                edgelines=el, \
                                sum_symbol=symb)
        types = np.array( [self._types for i in range(len(self))] )
        s2 = stringtools.string_summary( stringtools.strip_blank_lines( np.array2string(types, max_line_width=1000) ), \
                                threshold=th, \
                                edgelines=el, \
                                sum_symbol=symb)
        # replace single quotes by spaces:
        s2 = s2.replace("'", " ")
        return stringtools.two_columns( s1, s2, num_spaces=6, sep_symbol= '')
        
    def _show(self, show_entire=False):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        s = 'generations:\n'
        s += stringtools.indent( str(self.generations), 4) + '\n'
        s += 'array:\n'
        if show_entire:
            threshold = self._string_threshold
            self._string_threshold = np.nan
            s += stringtools.indent( self._data_labels2string(), 4 ) + '\n'
            self._string_threshold = threshold
        else:
            s += stringtools.indent( self._data_labels2string(), 4 ) + '\n'
        s += 'axes:\n'
        for i in range(self._ndim - 1):
            s += stringtools.indent( "%s  %s" % (i, self.labels['axes'][i]), 4 ) + '\n'
        return s
        
    def append(self, generation, data):
        if generation in self.generations:
            raise ValueError, "generation %d already stored" % generation
        self.generations.append( generation )
#        self._data_list.append( copy.deepcopy(np.array(data)) )
        self._data_list.append( np.array(data) )
        self._array_uptodate = False
    
    def _update_array(self):
        self.data = np.array( self._data_list )
        self._array_uptodate = True
    
    def get_sum(self, y):
        if y in self._axes:
            return self._get_axis_sum(y)
        try:
            return self._get_type_sum(y)
        except ValueError:
            raise ValueError, "'%s' is not a valid input. Input must either be an axis or type label." % str(y)
    
    def _get_axis_sum(self, _axis):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        result = condense(self.data, [0,self.get_position(_axis)+1])    # first axis is the generation axis!
        newlabels = dict(axes=[_axis], elements=[self._get_axis_elements(_axis)])
        code = "%s(generations=self.generations, data=result, labels=newlabels)" % self._return_class
        return eval(code)
    
    def _get_type_sum(self, _type):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        positions = []
        for element in _type.split('-'):
            positions.append(self.get_position(element))
        axis_positions = [0]                    # first axis is the generation axis!
        element_positions = [slice(None,None,None)]
        for ax,el in sorted(positions):
            axis_positions.append(ax+1)         # first axis is the generation axis!
            element_positions.append(el)
        return condense(self.data, axis_positions)[tuple(element_positions)]
        
    def _get_axis_slice(self, _axis):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        pos = self.get_position(_axis)
        result = np.sum(self.data, pos+1)
        newlabels = copy.deepcopy(self.labels)
        for x in ['axes', 'elements']:
            del newlabels[x][pos]
        code = "%s(generations=self.generations, data=result, labels=newlabels)" % self._return_class
        return eval(code)
    
    def _get_type_slice(self, _type):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        axis_labels = copy.copy( self.labels['axes'] )
        element_labels = copy.copy( self.labels['elements'] )
        positions = []
        for element in _type.split('-'):
            positions.append(self.get_position(element))
        slice_list = [slice(None,None,None) for i in range(self._ndim)]
        for ax,el in sorted(positions):
            slice_list[ax+1] = el            # first axis is the generation axis!
            ax_name = self.labels['axes'][ax]
            pos = axis_labels.index(ax_name)
            del axis_labels[pos]
            del element_labels[pos]
        result = self.data[ tuple(slice_list) ]
        labels = {'axes': axis_labels, 'elements': element_labels}
        code = "%s(generations=self.generations, data=result, labels=labels)" % self._return_class
        return eval(code)
        
    def getf(self, generation):
        # make sure that the data array is up to date:
        if not self._array_uptodate:
            self._update_array()
        pos = numpytools.get_closest_index( self.generations, generation )
        # print "returning frequencies for closest generation stored: %d" % self.generations[pos]
        return self[pos]
    
    def getg(self, generation):
        return numpytools.get_closest( self.generations, generation )

class Population(Frequencies):
    """Base class for all populations in specific models.
    
    Child classes should provide the appropriate data for the following class variables:
    _constant_reproduction_weights: []
    _dynamic_reproduction_weights: []
    _constant_viability_selection_weights: []
    """
    _return_class = 'Frequencies'
    
    def __init__(self, \
                 frequencies=np.array([[[1.,0.],[0.,0.]],[[0.,0.],[0.,0.]]]), \
                 labels={'axes': ['trait','preference','cytotype'], \
                         'elements': [['T1','T2'],['P0','P1'],['uninfected','Wolbachia']]}, \
                 generation=0):
        """weights={'constant reproduction': {'F': Weight, 'T': Weight, ...},
                    'dynamic reproduction':  {'SSW': Weight, ...},
                    'constant viability selection': {'S': Weight, ...}, 
                    'dynamic viability selection': {'?': Weight, ...}}
        """
        Frequencies.__init__(self, data=frequencies, labels=labels)
        self.generation = generation
        self._sex_mode = 'implicit'
        self._viability_selection_mode = 'sex independent'
        if labels['axes'][0] == 'sex':
            self._sex_mode = 'explicit'
            self._reproduction_axes = reproduction_axes(*labels['axes'][1:])
        elif labels['axes'][0] == 'y-chromosome':
            self._sex_mode = 'y-chromosome'
            self._reproduction_axes = reproduction_axes(*labels['axes'])
        else:
            self._reproduction_axes = reproduction_axes(*labels['axes'])
        self._reproduction_rank = len(self._reproduction_axes)
        self._previous_frequencies = np.zeros(self.data.shape)
        self._changes = np.zeros(self.data.shape)
        self._previous_changes = (-1) * np.ones(self.data.shape)
        self.eq_precision = 1e-7
        self.eq = self._eq_reached(self.eq_precision)
        self.history = FrequenciesHistory(generations=generation, data=self.data, labels=self.labels)
        
        self.weights = {}
        self._number_of_weights = 0
        for t in ['constant reproduction', 'dynamic reproduction', \
                  'constant viability selection']:
            self.weights[t] = {}
        
        #self._bind_methods()
    
    def __getstate__(self):
        result = self.__dict__.copy()
        try:
            del result['_reproduction']
        except:
            pass
        return result
    
    def __setstate(self, dict):
        self.__dict__ = dict
        #self._bind_methods()
    
    def show(self, generation=None, show_sums=True, show_entire=False, show_most_frequent=True):
        """If 'generation' is in ['parameters','weights','overview'] then a shortcut to the correspondign show method is called. Otherwise the _show method is called"""
        if generation in ['parameters','weights','overview','full overview']:
            code = "self.show_%s()" % generation.replace(' ', '_')
            exec code
            return
#        if generation == 'parameters':
#            self.show_parameters()
#            return
#        elif generation == 'weights':
#            self.show_weights()
#            return
#        elif generation == 'overview'):
#            self.show_overview()
#            return
        else:
            print self._show(generation, show_sums, show_entire, show_most_frequent)
    
    def show_weights(self):
#        print repr(self)
#        for _type in ['constant reproduction', 'dynamic reproduction', \
#                       'constant viability selection']:
#            print _type.swapcase() + ':'
#            if self.weights[_type] == {}:
#                print 'no weights of this type\n'
#            for name, w in self.weights[_type].iteritems():
#                print stringtools.two_columns(name,w._show())
        print self._show_weights()
    
    def show_parameters(self):
#        print repr(self)
#        for _type in ['constant reproduction', 'dynamic reproduction', \
#                       'constant viability selection']:
#            for w in self.weights[_type].values():
#                for pname in w._parameter_names:
#                    prettyparstring = pprint.pformat(w.__dict__[pname]).replace('{',' ').replace('_',' ').replace('}','').replace(',','')
#                    print stringtools.strip_blank_lines( stringtools.two_columns(pname+':', prettyparstring) )
##                    print "%s  %s\n" % (pname, pprint.pformat(w.__dict__[pname]))
        print self._show_parameters()

    
    def _show(self, mode=None, show_sums=True, show_entire=False, show_most_frequent=True):
        if mode in ['parameters','weights','overview','full overview']:
            code = "s = self._show_%s()" % mode.replace(' ', '_')
            exec code
            return s
        s = repr(self) + '\ngeneration: '
        if mode == None or mode == self.generation:
            s += str(self.generation) + '\n'
            f = Frequencies(self.data, self.labels)
        else:
            generation = mode
            f = Frequencies(self.history.getf(generation), self.labels)
            s += str(self.history.getg(generation)) + '\n'
        s += f._show(show_sums, show_entire, show_most_frequent)
        return s
    
    def _show_weights(self):
        s = repr(self) + '\n'
        for _type in ['constant reproduction', 'dynamic reproduction', \
                       'constant viability selection']:
            s += _type.swapcase() + ':\n'
            if self.weights[_type] == {}:
                s += 'no weights of this type\n'
            else:
                for name, w in self.weights[_type].iteritems():
                    s+= stringtools.two_columns(name,w._show()) + '\n'
        return s.strip()
    
    def _show_parameters(self):
        s = repr(self) + '\n'
        for _type in ['constant reproduction', 'dynamic reproduction', \
                       'constant viability selection']:
            for w in self.weights[_type].values():
                for pname in w._parameter_names:
                    prettyparstring = pprint.pformat(w.__dict__[pname]).replace('{',' ').replace('_',' ').replace('}','').replace(',','')
                    s+= stringtools.strip_blank_lines( stringtools.two_columns(pname+':', prettyparstring) ) + '\n'
        return s.strip()
        
#    def visualize(self, loci='all', generation=-1, withlabels=True, **kwargs):
#        f = plt.figure()
#        ax = plt.subplot(111)
#        animationtools.popviz(ax, self, generation=generation, withlabels=withlabels, **kwargs)
#        plt.draw()

    def _get_max_change(self):
        return np.max( self._get_first_order_changes() )
    
    def _state2string(self, *types):
        max_change = self._get_max_change()
        total_change = np.sum(self._changes)
        if len(types) == 1:
            freqs = [self.get_sum(types[0])]   # if len=1 we still need to extract the type from the tuple
        else:
            freqs = [self.get_sum(t) for t in types]
        s = "%10.3g%10.3g" % (total_change, max_change)
        for f in freqs:
            s += "%10.3g" % f
        return s
    
    def _shortstate2string(self, *types):
        if len(types) == 1:
            freqs = [self.get_sum(types[0])]   # if len=1 we still need to extract the type from the tuple
        else:
            freqs = [self.get_sum(t) for t in types]
        for f in freqs:
            s += "%10.3g" % f
        return s
               
    def _outputline(self, *types):
        print '\r',
        s = "%10d%s" % (self.generation, self._state2string(*types))
        sys.stdout.write(s)
        sys.stdout.flush()
    
    def _short_outputline(self, *types):
        print '\r',
        s = self._shortstate2string(*types)
        sys.stdout.write(s)
        sys.stdout.flush()
    
    def _state_title(self, *types):
        s = "%10s%10s" % ('total','max.')
        for t in types:
            if len(t) > 5:
                t = t[:5] + '.'
            s += "%10s" % t
        return s
    
    def _state_short_title(self, *types):
        for t in types:
            if len(t) > 5:
                t = t[:5] + '.'
            s += "%10s" % t
        return s
    
    def _output_title(self, *types):
        s = "%10s" % 'generation'
        s += self._state_title(*types)
        return s
    
    def _short_output_title(self, *types):
        s  = "%10s" % 'generation'
        s += self._state_short_title(*types)
        return s
    
    def add_weight(self, weight, update_constant_factors=True):
        if weight._scope == 'population':
            weight._refer_to_pop(self)
        weight._init_after_assignment()
        wtype = weight._type
        if wtype not in self.weights.keys():
            raise TypeError, "type of weight '%s' is unknown" % wtype
        wname = "w%d" % self._number_of_weights
        self._number_of_weights += 1
        self.weights[wtype][wname] = weight
        #self._bind_methods()
        if update_constant_factors:
            self._build_constant_reproduction_factor()
            self._build_constant_viability_selection_factor()
    
    def add_weights(self, *weights):
        for w in weights:
            self.add_weight(w, update_constant_factors=False)
        self._build_constant_reproduction_factor()
        self._build_constant_viability_selection_factor()
    
    def change_weight(self, wname, weight, update_constant_factors=True):
        weight._refer_to_pop(self)
        weight._init_after_assignment()
        wtype = weight._type
        self.weights[wtype][wname] = weight
        #self._bind_methods()
        if update_constant_factors:
            self._build_constant_reproduction_factor()
            self._build_constant_viability_selection_factor()
    
    def _init_weights(self):
        # expand and form into one all weights that are constant:
        self._build_constant_reproduction_factor()
        self._build_constant_viability_selection_factor()
    
    def _bind_methods(self):
        # bind _reproduction to the appropriate method:
        if self.weights['dynamic reproduction']:
            if self._sex_mode == 'implicit':
                self._reproduction = self._reproduction_dynamic_implicit_sex
            elif self._sex_mode == 'explicit':
                self._reproduction = self._reproduction_dynamic_explicit_sex
            elif self._sex_mode == 'y-chromosome':
                self._reproduction = self._reproduction_dynamic_y_chromosome
        else:
            if self._sex_mode == 'implicit':
                self._reproduction = self._reproduction_constant_implicit_sex
            elif self._sex_mode == 'explicit':
                self._reproduction = self._reproduction_constant_explicit_sex
            elif self._sex_mode == 'y-chromosome':
                self._reproduction = self._reproduction_constant_y_chromosome
        # bind _viability_selection to the appropriate method:
        if self._viability_selection_mode == 'sex dependent':
            if self._sex_mode == 'explicit':
                self._viability_selection = self._viability_selection_sex_dependent
        elif self._viability_selection_mode == 'sex independent':
            self._viability_selection = self._viability_selection_sex_independent
    
    def _detach_methods(self):
        self._reproduction = None
        self._viability_selection = None
    
    def _build_constant_reproduction_factor(self):
        result = np.array(1.)
        for w in self.weights['constant reproduction'].values():
            result = result * w.expanded(self._reproduction_axes)
#        self._constant_R_factor = copy.deepcopy(result)
        self._constant_R_factor = result
    
    def _build_dynamic_reproduction_factor(self):
        result = 1.
        for w in self.weights['dynamic reproduction'].values():
            result = result * w.expanded( self._reproduction_axes )
#        self._dynamic_R_factor = copy.deepcopy(result)
        self._dynamic_R_factor = result
        
    def _build_constant_viability_selection_factor(self):
        result = 1.
        if self.weights['constant viability selection'] != {}:
            for w in self.weights['constant viability selection'].values():
                result = result * w.expanded( self._axes )
#        self._constant_VS_factor = copy.deepcopy(result)
        self._constant_VS_factor = result
        
    def _mostfrequent2string(self, n=8, sum_axes=None):
        return Frequencies._mostfrequent2string(self,n,sum_axes)
            
    def _hstore(self):
        """Store current state in history."""
        if self.generation != self.history.generations[-1]:
            self.history.append(self.generation, self.data)
    
    def _threshold_hstore(self, threshold=1e-2):
        """Store current state if frequencies have changed by more than threshold
since the last storage."""
        # if change since last store event is larger than threshold then store current state:
        if self._get_total_change_since_last_hstore() >= threshold:
            self._hstore()
    
    def _get_first_order_changes(self):
        """Absolute differences between current and previous frequencies.
Returns an array of the same shape as data."""
        return np.abs( self.data - self._previous_frequencies )
    
    def _get_total_change_since_last_hstore(self):
        """Sum of absolute differences between current state frequencies and last stored.
Returns a single float."""
        return numpytools.absolute_difference_sum( self.data, self.history[-1] )

    def _get_changes_since_last_hstore(self):
        """Absolute differences between current state frequencies and last stored."""
        return np.abs( self.data - self.history[-1] )
    
    def _sex_normalize(self):
        """Assumes that modeling is sex explicit, and that first axis is for 'sex'; normalizes females and males independent from each other."""
        l = len(self._axes)
        norm = numpytools.condense( self.data, 0)   # sex specific norm
        self.data /= numpytools.expand(norm, l, 0)  
        self.normalize()      

    def _sex_normalize(self):
        """Assumes that modeling is either sex explicit or by y-chromosome, and that first axis is for 'sex' or 'y-chromosome', respectively; normalizes females and males independent from each other."""
        l = len(self._axes)
        y_sum = numpytools.sum_except_for( self.data, 0 )
        female_sum = y_sum[0]
        male_sum = np.sum( y_sum[1:] )
        norm = male_sum * np.ones(self._data_shape)
        norm[0,:] = female_sum
        self.data /= norm
        self.normalize()      
    
    def set_migration(self, *args):
        self.migration_rates = list(args[::2])
        self.source_frequencies = list(args[1::2])
        n = len(self.migration_rates)
        for i in range(n):
            self.__dict__["migrate%d" % (i+1)] = self.migration_rates[i]
            self.__dict__["mainland%d" % (i+1)] = Frequencies(self.source_frequencies[i], self.labels)
    
    def _migration(self):
        migratesum = 0.
        inputsum = 0.
        for m,f in zip(self.migration_rates, self.source_frequencies):
            migratesum += m
            inputsum += m*f
        self.data *= (1.-migratesum)
        self.data += inputsum
    
    def _viability_selection(self):
        print "dummy viability selection method"
        
    def _viability_selection_sex_independent(self):
        """Use this viability selection method if sexes are treated the same."""
        self.set_array( self.data * self._constant_VS_factor )
        self.normalize()

    def _viability_selection_sex_dependent(self):
        """Use this viability selection method if sexes are treated differently and stated explicitely or are defined by y-chromosome."""
        self.set_array( self.data * self._constant_VS_factor )
        #self._sex_normalize()
        self.normalize()
    
    def introduce_allele(self, allele, freq, advance=True):
        # check whether allele not already existent:
        if self.get_sum(allele) != 0.:
            raise KeyError, "allele %s already in population" % allele
        ax, el = self.get_position(allele)        # axis index, element index
        newfreqs = copy.deepcopy(self.data)
        slice_list = [slice(None,None,None) for i in range(self._ndim)]
        slice_list[ax] = el
        axis_name = self._get_element_axis(allele)
        newfreqs[slice_list] = self[axis_name]
        introfreqs = (1-freq) * np.ones((self._data_shape[ax]))
        introfreqs[el] = freq
        newfreqs *= numpytools.expand(introfreqs, self._ndim, ax)
        # bookkeeping and data setting:
        self._bookkeeping_before_cycle()
        self.set_array(newfreqs)
        self.normalize()
        self._bookkeeping_after_cycle()
        if advance:
            self.generation_increase(1)
            
    def _set_allele_frequency(self, allele, freq, advance=True):
        """Set the frequency of allele 'allele' to 'freq'. All other frequencies stay more or less the same."""
        ax, el = self.get_position(allele)        # axis index, element index
        newfreqs = copy.deepcopy(self.data)
        slice_index = [slice(None,None,None) for i in range(self._ndim)]
        slice_index[ax] = el
        allele_sum = self.get_sum(allele)
        if allele_sum == 0.:
            self.introduce_allele(allele, freq)
            return
        array_norm = np.ones(np.shape(newfreqs)) * (1-freq)/(1-allele_sum)
        array_norm[slice_index] = freq/allele_sum
        print 'norm:\n', array_norm
        newfreqs *= array_norm
        if advance:
            # bookkeeping, set data, and normalize:
            self._bookkeeping_before_cycle()
            self.set_array(newfreqs)
            self.normalize()
            self._bookkeeping_after_cycle()
            self.generation_increase(1)
        else:
            # only set data and normalize:
            self.set_array(newfreqs)
            self.normalize()
        
    def _set_type_frequency(self, _type, freq, advance=False):
        """Set the frequency of the complete nucleocytotype '_type' to 'freq'.
        """
        newdata = copy.copy(self.data)
        pos = self._get_index(_type)
        newdata[pos] = 0.
        newdata *= (1.-freq)/np.sum(newdata)
        newdata[pos] = freq
        self.set_array(newdata)
        self.normalize()
            
    def _reproduction_constant_implicit_sex(self):
        """Reproduction method for the case WITHOUT dynamic reproduction weights."""
        r = len(self._reproduction_axes)
        females = expand( self.data, r, range(0,r/3) )
        males = expand( self.data, r, range(r/3,2*r/3) )
        offspring = condense( females * males * \
                              self._constant_R_factor, \
                    range(2*r/3,r) )
#        pdb.set_trace()
        self.set_array(offspring)
        self.normalize()
    
    def _reproduction_dynamic_implicit_sex(self):
        """Reproduction method for the case WITH dynamic reproduction weights."""
#        pdb.set_trace()
        r = len(self._reproduction_axes)
        self._update_dynamic_reproduction_weights()
        self._build_dynamic_reproduction_factor()
        females = expand( self.data, r, range(0,r/3) )
        males = expand( self.data, r, range(r/3,2*r/3) )
        offspring = condense( females * males * \
                              self._constant_R_factor * \
                              self._dynamic_R_factor, \
                    range(2*r/3,r) )
        self.set_array(offspring)
        self.normalize()
    
    def _update_dynamic_reproduction_weights(self):
        for w in self.weights['dynamic reproduction'].values():
            w.update()
    
    def _reproduction_constant_explicit_sex(self):
        """Reproduction method for the case WITHOUT dynamic reproduction weights."""
        r = len(self._reproduction_axes)
        females = expand( self.data[0], r, range(0,r/3) )
        males = expand( self.data[1], r, range(r/3,2*r/3) )
        offspring = condense( females * males * \
                              self._constant_R_factor, \
                    range(2*r/3,r) )
        self.set_array( [0.5*offspring, 0.5*offspring] )        # 1:1 sex ratio
        self.normalize()
    
    def _reproduction_constant_y_chromosome(self):
        """Reproduction method for the case WITHOUT dynamic reproduction weights."""
        r = len(self._reproduction_axes)
        females = expand( self.data, r, range(0,r/3) )
        males = expand( self.data, r, range(r/3,2*r/3) )
        offspring = condense( females * males * \
                              self._constant_R_factor, \
                    range(2*r/3,r) )
        self.set_array( offspring )
        self.normalize()      # females and males sum up to 1. each

    def _reproduction_dynamic_explicit_sex(self):
        """Reproduction method for the case WITH dynamic reproduction weights."""
        r = len(self._reproduction_axes)
        #for wname in self._dynamic_reproduction_weight_names:
        #    self.weights['dynamic reproduction'][wname].update()
        self._update_dynamic_reproduction_weights()
        self._build_dynamic_reproduction_factor()
        females = expand( self.data[0], r, range(0,r/3) )
        males   = expand( self.data[1], r, range(r/3,2*r/3) )
        offspring = condense( females * males * \
                              self._constant_R_factor * \
                              self._dynamic_R_factor, \
                    range(2*r/3,r) )
        self.set_array( [0.5*offspring, 0.5*offspring] )        # 1:1 sex ratio
        self.normalize()
    
    def _reproduction_dynamic_y_chromosome(self):
        """Reproduction method for the case WITH dynamic reproduction weights."""
        r = len(self._reproduction_axes)
        #for wname in self._dynamic_reproduction_weight_names:
        #    self.weights['dynamic reproduction'][wname].update()
        self._update_dynamic_reproduction_weights()
        self._build_dynamic_reproduction_factor()
        females = expand( self.data, r, range(0,r/3) )
        males   = expand( self.data, r, range(r/3,2*r/3) )
#        print 'females   ', np.shape(females)
#        print 'males     ', np.shape(males)
#        print 'constant R', np.shape(self._constant_R_factor)
#        print 'dynamic R ', np.shape(self._dynamic_R_factor)
        offspring = condense( females * males * \
                              self._constant_R_factor * \
                              self._dynamic_R_factor, \
                    range(2*r/3,r) )
        self.set_array( offspring )
        self.normalize()      # females and males sum up to 1. each
    
#    def _set_repro_func(self):
#        if self.weights['dynamic reproduction']:
#            if self._sex_mode == 'implicit':
#                reproduction = self._reproduction_dynamic_implicit_sex
#            elif self._sex_mode == 'explicit':
#                reproduction = self._reproduction_dynamic_explicit_sex
#        else:
#            if self._sex_mode == 'implicit':
#                reproduction = self._reproduction_constant_implicit_sex
#            elif self._sex_mode == 'explicit':
#                reproduction = self._reproduction_constant_explicit_sex
#        self._reproduction = reproduction
    
    def _bookkeeping_before_cycle(self):
        self._previous_frequencies = copy.deepcopy(self.data)
        self._previous_changes = copy.deepcopy(self._changes)
    
    def _bookkeeping_after_cycle(self):
        self._changes = self._get_first_order_changes()
    
    def _reproduction(self):
        print 'dummy reproduction method'
        
    def _next_generation(self):
        """Warning: You must have set the reproduction and viability selection methods via 
        _bind_methods prior to using this method."""
        self._bookkeeping_before_cycle()
        # cycle:
        self._migration()
        self._viability_selection()
        self._reproduction()
        self.generation += 1
        self._bookkeeping_after_cycle()
        
    def _next_generation_isolated(self):
        """Warning: You must have set the reproduction and viability selection methods via 
        _bind_methods prior to using this method."""
        self._bookkeeping_before_cycle()
        # cycle:
        self._viability_selection()
        self._reproduction()
        self.generation += 1
        self._bookkeeping_after_cycle()
    
    def _next_generation_within_metapop(self):
        """Warning: You must have set the reproduction and viability selection methods via 
        _bind_methods prior to using this method."""
        # the metapopulation handles the bookkeeping
        # cycle:
        self._viability_selection()
        self._reproduction()
        self.generation += 1
    
    def run(self, n, isolated=False, output_modulo=20, output_types=None):
        """Run for n generations."""
        self._build_constant_reproduction_factor()
        self._build_constant_viability_selection_factor()
        self._bind_methods()
        if isolated:
            try:
                if output_types:
                    print self._output_title(*output_types)
                    for i in range(n):
                        self._next_generation_isolated()
                        if self.generation % output_modulo == 0:
                            self._outputline(*output_types)
                        self._threshold_hstore()
                        numpytools.ipbreak()
                else:
                    for i in range(n):
                        self._next_generation_isolated()
                        self._threshold_hstore()
                        numpytools.ipbreak()
            finally:
                self._hstore()
                #self._detach_methods()
        else:
            try:
                if output_types:
                    print self._output_title(*output_types)
                    for i in range(n):
                        self._next_generation()
                        if self.generation % output_modulo == 0:
                            self._outputline(*output_types)
                        self._threshold_hstore()
                        numpytools.ipbreak()
                else:
                    for i in range(n):
                        self._next_generation()
                        self._threshold_hstore()
                        numpytools.ipbreak()
            finally:
                self._hstore()
                self._detach_methods()
    
    def runto(self, nexact, isolated=False, output_modulo=20, output_types=None):
        """Run until generation n is reached."""
        n = nexact - self.generation
        self.run(n, isolated=isolated, output_modulo=output_modulo, output_types=output_types)
    
    def _first_order_changes_all_smaller_than(self, precision):
        return np.alltrue( self._changes < precision )
    
    def _second_order_changes_all_decreasing(self):
        return np.alltrue( self._previous_changes >= self._changes )
    
    def _total_changes_decreasing(self):
        return np.sum(self._previous_changes) >= np.sum(self._changes)
    
    def reset(self):
        self.generation = 0
        self.set_array(self.history[0])
        self._changes = np.zeros(self.data.shape)
        self._previous_changes = (-1.) * np.ones(self.data.shape)
        self.history = FrequenciesHistory(generations=self.generation, data=self.data, labels=self.labels)
        
    def _eq_reached(self, precision=1e-7):
        if not self._first_order_changes_all_smaller_than(precision):
            return False
        else:
            #if not self._second_order_changes_all_decreasing():
            if not self._total_changes_decreasing():
                return False
            else:
                return True
    
    def _allele_increases(self, allele='P1'):
        previous_freqs = Frequencies(self._previous_frequencies, self.labels)
#        print 'previous frequencies:', previous_freqs
        current = self.get_sum(allele)
        previous = previous_freqs.get_sum(allele)
        return current > previous
            
#    def _eq_reached(self, precision=1e-7):
#        if not self._first_order_changes_all_smaller_than(precision):
#            return False
#        else:
#            if (np.sum(self._previous_changes) > np.sum(self._changes)):
#                return False
#            else:
#                return True

    def findEQ(self, precision=1e-7, nmin=1, nmax=5e5, isolated=False, interval=100, output_types=None):
        """Find EQ frequencies.
Run until all individual frequencies do not change more than precision between 
consecutive generations. Run for at least nmin generations, and stop simulation
after nmax generation even if no EQ has been reached.
Divide run into parts of 'interval' generations to avoid eq checking overhead.
"""
        self.eq_precision = precision
        self.eq = self._eq_reached(precision)
        if self.eq:
            return True
        self._build_constant_reproduction_factor()
        self._build_constant_viability_selection_factor()
        self._bind_methods()
        nmax += self.generation
        if isolated:
            self.run(nmin, isolated=True)
            try:
                if output_types:
                    print
                    print self._output_title(*output_types)
                    while not self._eq_reached(precision) and \
                              self.generation < nmax:
                        for i in range(interval):
                            self._next_generation_isolated()
                            self._threshold_hstore()
                            numpytools.ipbreak()
                        self._outputline(*output_types)
                else:
                    while not self._eq_reached(precision) and \
                              self.generation < nmax:
                        for i in range(interval):
                            self._next_generation_isolated()
                            self._threshold_hstore()
                            numpytools.ipbreak()
            finally:
                if output_types:
                    print
                self._hstore()
                self.eq = self._eq_reached(precision)
                return self.eq
        else:
            self.run(nmin)
            try:
                if output_types:
                    print
                    print self._output_title(*output_types)
                    while not self._eq_reached(precision) and \
                              self.generation < nmax:
                        for i in range(interval):
                            self._next_generation()
                            self._threshold_hstore()
                            numpytools.ipbreak()
                        self._outputline(*output_types)
                else:
                    while not self._eq_reached(precision) and \
                              self.generation < nmax:
                        for i in range(interval):
                            self._next_generation()
                            self._threshold_hstore()
                            numpytools.ipbreak()
            finally:
                if output_types:
                    print
                self._hstore()
                self.eq = self._eq_reached(precision)
                return self.eq
    
    def findIncrease(self, allele='P1', precision=1e-7, nmin=10, repetition=5, isolated=False, output_types=None):
        """Find if spread of 'allele' occurs.
Run for at 'nmin' generations. Returns True if frequency of 'allele'
increases for 'repetition' successive generations.
"""
        self._build_constant_reproduction_factor()
        self._build_constant_viability_selection_factor()
        self._bind_methods()
        increases = []
        if isolated:
            self.run(nmin, isolated=True)
            try:
                if output_types:
                    print
                    print self._output_title(*output_types)
                    for i in range(repetition):
                        self._next_generation_isolated()
                        self._hstore()
                        increases.append( self._allele_increases(allele) )
                        numpytools.ipbreak()
                    self._outputline(*output_types)
                else:
                    for i in range(repetition):
                        self._next_generation_isolated()
                        self._hstore()
                        increases.append( self._allele_increases(allele) )
                        numpytools.ipbreak()
            finally:
                if output_types:
                    print
                self._hstore()
                return np.alltrue( increases )
        else:
            self.run(nmin)
            try:
                if output_types:
                    print
                    print self._output_title(*output_types)
                    for i in range(repetition):
                        self._next_generation()
                        self._hstore()
                        increases.append( self._allele_increases(allele) )
                        numpytools.ipbreak()
                    self._outputline(*output_types)
                else:
                    for i in range(interval):
                        self._next_generation()
                        self._hstore()
                        increases.append( self._allele_increases(allele) )
                        numpytools.ipbreak()
            finally:
                if output_types:
                    print
                self._hstore()
                return np.alltrue( increases )
    
    def generation_increase(self, inc=1):
        """Increase the generation count by 1 and store history."""
        self.generation += inc
        self._hstore()
    
    def get_frequencies(self):
        return Frequencies(self.data, self.labels)
    
    def get_weights(self):
        """Returns an (unsorted) list of all population weights."""
        weights = []
        for t in self.weights.keys():
            for w in self.weights[t].values():
                weights.append(w)
        return weights

class FocalPopulation(Population):
    """This class assumes that a neutral locus 'neutral' exists at the last axis with alleles 'np1','np2','np3'.
    
    Calling focal_normalization() distributes the frequency of allele 'np3' to the other two alleles 'np1' and 'np2', at the same time ensuring that the sums at the other axes do not change."""
    def focal_normalization(self):
        d = self._ndim

        n1 = self['np1'].data
        n2 = self['np2'].data
        n3 = self['np3'].data

        newfreqs = np.zeros(self._data_shape)
        
        n1new = np.nan_to_num( n1 + n3 * n1/(n1+n2) )            # distribute n3 to n1 and n2 with existing ratio and replace occurences of nan with zeros
        ax,el = self.get_position('np1')
        slice_index = [slice(None,None,None) for i in range(d)]
        slice_index[ax] = el
        newfreqs[slice_index] = n1new

        n2new = np.nan_to_num( n2 + n3 * n2/(n1+n2) )            # distribute n3 to n1 and n2 with existing ratio and replace occurences of nan with zeros
        ax,el = self.get_position('np2')
        slice_index = [slice(None,None,None) for i in range(d)]
        slice_index[ax] = el
        newfreqs[slice_index] = n2new
        
        # just for completeness: n3new is zero
        # n3new = 0.
        
        self.set_array(newfreqs)
    
    def _next_generation_within_metapop(self):
        """FocalPopulation overwrites this method to include the focal normalization."""
        # the metapopulation handles the bookkeeping
        # cycle:        
        self.focal_normalization()
        self._viability_selection()
        self._reproduction()
        self.generation += 1

    
class BackgroundPopulation(Population):
    """This class assumes that a neutral locus 'neutral' exists at the last axis with alleles 'np1','np2','np3'.
    
    Calling background_normalization() sets the frequency of allele 'np3' to 1. It is ensured that the sums at all other axes do not change."""
    def background_normalization(self):
        d = self._ndim
        newfreqs = np.zeros(self._data_shape)    
        ax,el = self.get_position('np3')
        slice_index = [slice(None,None,None) for i in range(d)]
        slice_index[ax] = el
        newfreqs[slice_index] = 1.
        
        positions = range(d)
        del positions[ax]
        norm = self['neutral'].data
        
        self.set_array( newfreqs * numpytools.expand(norm, d, positions) )
    
    def _next_generation_within_metapop(self):
        """BackgroundPopulation overwrites this method to include the background normalization."""
        # the metapopulation handles the bookkeeping
        # cycle:
        self.background_normalization()
        self._viability_selection()
        self._reproduction()
        self.generation += 1
    
        

def main():
    import weights4 as weights
    reload(weights)
    freqs = np.zeros((2,2,2),float)
    freqs[0,0,0] = 1.
    labels = dict(axes=['trait','preference','cytotype'], \
                  elements=[['T1','T2'],['P0','P1'],['uninfected','Wolbachia']])
    pop = Population(freqs,labels)
    CI = weights.CytoplasmicIncompatibility(0.9)
    F = weights.FecundityReduction(0.)
    NITP = weights.NuclearInheritanceAtTwoLoci(0.5, 'trait', 'T', range(1,3), 'preference', 'P', range(2))
    MP = weights.MatingPreference(0.99, 'pr', 'P0','None',0., 'P1','T1',0.99)
    VS = weights.ViabilitySelection(selection_coefficients=[1.,0.])
    T = weights.WolbachiaTransmission(0.87)
    pop.add_weights(CI,F,NITP,MP,VS,T)
    #pop.findEQ(isolated=True)
    #pop.show()
    #pop.introduce_allele('Wolbachia', 0.1)
    #pop.show()
    #pop.findEQ(isolated=True)
    #pop.show()
    return pop

